package com.wyr.leetcode.step2.dfs;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 给你二叉搜索树的根节点 root ，该树中的 恰好 两个节点的值被错误地交换。请在不改变其结构的情况下，恢复这棵树 。
 *
 * https://leetcode.cn/problems/recover-binary-search-tree/description/?envType=problem-list-v2&envId=depth-first-search
 */
public class RecoverTree {

    public static void main(String[] args) {
        TreeNode node1 = new TreeNode(1);
        TreeNode node2 = new TreeNode(2);
        TreeNode node3 = new TreeNode(3);
        TreeNode node4 = new TreeNode(4);

        node3.left=node1;
        node3.right=node4;
        node4.left=node2;

        node3.right = new TreeNode(2);

        recoverTree(node3);
    }

    public static void recoverTree(TreeNode root) {
        // 二叉搜索树上的节点的值不会有重复的值
        // 记录中序遍历经过的值的路径
        List<Integer> midTrace = new ArrayList<>();
        // 保存值和节点的映射关系
        Map<Integer, TreeNode> map = new HashMap<>();
        // 填充 midTrace和map
        dfs(root,midTrace, map);
        // 找到有问题的节点的值
        List<Integer> errorNodesVal = findErrorNodes(midTrace);
        // 进行节点值的交换
        swapErrorNodeVal(errorNodesVal, map);
    }

    private static void swapErrorNodeVal(List<Integer> errorNodes, Map<Integer, TreeNode> map) {
        int tmp = map.get(errorNodes.get(0)).val;
        map.get(errorNodes.get(0)).val = map.get(errorNodes.get(1)).val;
        map.get(errorNodes.get(1)).val = tmp;
    }

    public static void dfs(TreeNode root, List<Integer> midTrace, Map<Integer, TreeNode> map){
        if(root == null){
            return;
        }
        dfs(root.left, midTrace, map);
        midTrace.add(root.val);
        map.put(root.val, root);
        dfs(root.right, midTrace, map);
    }


    public static List<Integer> findErrorNodes(List<Integer> midTrace) {
        List<Integer> res = new ArrayList<>();
        for (int i = 0; i < midTrace.size(); i++) {
            int val = midTrace.get(i);
            if(midTrace.get(i+1)<val){
                // 这个值肯定是有问题的
                res.add(val);
                for (int j = i+2; j < midTrace.size(); j++){
                    if(midTrace.get(j)< midTrace.get(i+1)){
                        res.add(midTrace.get(j));
                        break;
                    }
                }
                // 没找到，就加入i+1处的元素
                if(res.size() == 1){
                    res.add(midTrace.get(i+1));
                    break;
                }
            }
            if (!res.isEmpty()){
                break;
            }
        }
        return res;
    }
}
